package org.apache.spark.examples

import org.apache.spark.{SparkConf, SparkContext}
import scala.util.Random

object RDDGroupByShuffleOperations {
  def main(args: Array[String]): Unit = {
    // 设置更多的分区以便观察shuffle
    val conf = new SparkConf()
      .setAppName("RDD GroupBy Shuffle Operations")
      .setMaster("local[4]")  // 使用4个核心以观察并行处理
    val sc = new SparkContext(conf)

    // 1. 创建大量随机数据，使用多个分区
    // stage 0
    val random = new Random(42) // 使用固定种子以保证结果可重现
    val numbersRDD = sc.parallelize(1 to 1000000, 8) // 100万个数字，8个分区
      .map(_ => random.nextInt(100)) // 生成0-99的随机数
    
    println("\n=== 原始数据分区信息 ===")
    println(s"分区数量: ${numbersRDD.getNumPartitions}")
    println(s"是否存在shuffle: 否（还未进行groupBy操作）")

    // 2. 执行groupBy操作，按照数字范围分组
    println("\n=== 执行groupBy操作 ===")
    println("注意：这将触发shuffle操作")

    // stage 1
    val rangeGroups = numbersRDD.groupBy(num => {
      if (num < 20) "0-19"
      else if (num < 40) "20-39"
      else if (num < 60) "40-59"
      else if (num < 80) "60-79"
      else "80-99"
    })

    // 3. 计算每个分组的统计信息
    val stats = rangeGroups.mapValues(values => {
      val count = values.size
      val sum = values.sum
      val avg = sum.toDouble / count
      (count, sum, avg)
    })

    // 4. 强制执行计算以显示结果和触发shuffle
    println("\n=== 分组统计结果 ===")
    stats.collect().toSeq.sortBy(_._1).foreach { case (range, (count, sum, avg)) =>
      println(f"范围 $range: 数量=$count, 总和=$sum, 平均值=$avg%.2f")
    }

    // 5. 执行另一个会产生shuffle的操作：按照是否为质数分组
    println("\n=== 执行质数分组（另一个shuffle操作）===")
    
    def isPrime(n: Int): Boolean = {
      if (n <= 1) return false
      if (n == 2) return true
      !(2 to math.sqrt(n).toInt).exists(x => n % x == 0)
    }

    val primeGroups = numbersRDD.groupBy(num => isPrime(num))
    
    val primeStats = primeGroups.mapValues(values => {
      val count = values.size
      val avg = values.sum.toDouble / count
      (count, avg)
    })

    println("\n=== 质数分组统计结果 ===")
    primeStats.collect().foreach { case (isPrime, (count, avg)) =>
      println(f"是否为质数=$isPrime: 数量=$count, 平均值=$avg%.2f")
    }

    // 6. 使用不同分区数执行groupBy以观察shuffle的影响
    println("\n=== 使用不同分区数的groupBy ===")
    
    val numbers2RDD = numbersRDD.repartition(4) // 重新分区为4个
    val rangeGroups2 = numbers2RDD.groupBy(num => num % 10) // 按个位数分组
    
    println(s"重新分区后的分区数: ${rangeGroups2.getNumPartitions}")
    
    val stats2 = rangeGroups2.mapValues(values => {
      val count = values.size
      val avg = values.sum.toDouble / count
      (count, avg)
    })

    println("\n=== 个位数分组统计结果 ===")
    stats2.collect().toSeq.sortBy(_._1).foreach { case (digit, (count, avg)) =>
      println(f"个位数=$digit: 数量=$count, 平均值=$avg%.2f")
    }

//    println("\n现在您可以在Spark UI中观察shuffle操作的详细信息：")
//    println("1. 访问 http://localhost:4040")
//    println("2. 查看Stages标签页")
//    println("3. 观察每个stage的shuffle读写数据量")
//    println("4. 查看DAG可视化图了解shuffle依赖")

    // 暂停足够长的时间以便查看Spark UI
    println("\n程序将暂停5分钟，请在此期间查看Spark UI...")
    Thread.sleep(300000) // 暂停5分钟

    sc.stop()
  }
} 